#include "comm/znx_string.h"
#include "comm/znx_inet.h"
#include "comm/znx_random.h"
#include "event/znx_conn.h"
#include "event/znx_tcp_proc.h"
#include "event/znx_udp_proc.h"
#include "event/znx_worker_proc.h"
#include "event/znx_session_pool.h"

// Five scenarios were tested
// 1. The idle connection exceeds the upper limit of the connection pool.
// 2. The idle connection exceeds the upper limit of a single host in the connection pool
// 3. The connection put into the connection pool was closed due to timeout
// 4. The server waits for a timeout and actively shuts down
// 5. proc_exit, close all idle connections.


static void
znx_test_server_read(znx_event_t *event)
{
    znx_conn_t *server_conn = event->data;
    char data[512];
    size_t data_len;
    znx_syscall_result_t res;

    if (event->timedout) {
        ZNX_DEF_LOG_DEBUG("test_server: read timeout");
        znx_event_loop_remove_conn(server_conn);
        znx_conn_close(server_conn);
        znx_thread_free(server_conn);
        return;
    }

    while (1) {
        res = znx_conn_read(server_conn, data, 500);
        if (res.res > 0) {
            data[res.res] = '\0';
            data_len = (size_t)res.res;
            ZNX_DEF_LOG_DEBUG("test_server read data: %s", data);
            continue;
        }

        if (res.eno == EAGAIN) {
            ZNX_DEF_LOG_DEBUG("test_server read EAGAIN");
            break;
        }

        ZNX_DEF_LOG_DEBUG("test_server: client close");
        znx_event_loop_remove_conn(server_conn);
        znx_conn_close(server_conn);
        znx_thread_free(server_conn);
        return;
    }

    res = znx_conn_write(server_conn, data, data_len);
    if ((size_t)res.res != data_len) {
        ZNX_DEF_LOG_FATAL("send failed, expect %d, but got %d", data_len, res.res);
    }

    int64_t r = znx_random_n(2);
    int64_t timeout_ms = 100;
    if (r == 1) {
        timeout_ms = 500;
    }

    znx_event_loop_add_timer(&server_conn->read, timeout_ms);
}


static void
tcp_server_accept_handler(znx_event_t *event)
{
    znx_conn_t *server_conn = event->data;

    server_conn->read.handler = znx_test_server_read;

    znx_event_loop_start_io(server_conn);
    // enabled read
    znx_event_loop_enable_read(server_conn);

    // znx_test_server_read(event);
}


static void client_conn_free(znx_conn_t *conn)
{
    if (conn->read.timedout) {
        ZNX_DEF_LOG_DEBUG("client conn %p free callback, timedout", conn);
    } else {
        char buf[4];
        znx_syscall_result_t sres;
        sres = znx_conn_read(conn, buf, 4);
        if (sres.res == 0) {
            ZNX_DEF_LOG_DEBUG("client conn %p free callback, server close", conn);
        } else if (sres.eno != 0 && sres.eno != EAGAIN) {
            ZNX_DEF_LOG_DEBUG("client conn %p free callback, errno: %d", conn, sres.eno);
        }
    }

    ZNX_DEF_LOG_DEBUG("client_conn close");
    znx_event_loop_remove_conn(conn);
    znx_conn_close(conn);
    znx_thread_free(conn);
}


static void test_get_and_put(znx_str_t *address_str)
{
    znx_address_t   address;
    int             port = 0;
    bzero(&address, sizeof(znx_address_t));
    znx_address_parser(address_str->data, address_str->len, &address, &port);

    znx_conn_t *conn =
        znx_g_session_pool_get(address_str, &address);
    if (conn == NULL) {
        ZNX_DEF_LOG_DEBUG("client session get %V failed", address_str);
        return;
    } else {
        ZNX_DEF_LOG_DEBUG("client session get %V success, conn: %p", address_str, conn);
    }

    int timeout_ms = 50;
    int64_t r = znx_random_n(2);
    if (r == 1) {
        timeout_ms = 200;
    }

    znx_str_t *peer_address_text =
        znx_conn_peer_address_text(conn);

    znx_g_session_pool_put(peer_address_text, conn, timeout_ms, client_conn_free);

    ZNX_DEF_LOG_DEBUG("session_pool_put: %p, address: %V",
        conn, address_str);
}


static void client_timer_handler(znx_event_t *event)
{
    znx_thread_free(event);
    znx_str_t address;
    
    ZNX_STR_SET(&address, "127.0.0.1:8888");
    test_get_and_put(&address);
    ZNX_STR_SET(&address, "127.0.0.1:8889");
    test_get_and_put(&address);
    ZNX_STR_SET(&address, "127.0.0.1:8889");
    test_get_and_put(&address);
    ZNX_STR_SET(&address, "127.0.0.1:8888");
    test_get_and_put(&address);
    ZNX_STR_SET(&address, "127.0.0.1:8888");
    test_get_and_put(&address);
    ZNX_STR_SET(&address, "127.0.0.1:8889");
    test_get_and_put(&address);
}


static void client_read(znx_event_t *event)
{
    znx_conn_t *conn;
    conn = event->data;
    znx_syscall_result_t res;
    char buf[512];

    // int id = (int)(intptr_t)conn->data;

    while (1) {
        res = znx_conn_read(conn, buf, 512);
        if (res.res > 0) {
            buf[res.res] = '\0';
            ZNX_DEF_LOG_DEBUG("test_client read data: %s", buf);
            continue;
        }

        if (res.eno == EAGAIN) {
            break;
        }

        ZNX_DEF_LOG_DEBUG("test_client read failed");
        znx_event_loop_remove_conn(conn);
        znx_conn_close(conn);
        znx_thread_free(conn);
    }

    znx_str_t *peer_address_text =
        znx_conn_peer_address_text(conn);
    znx_g_session_pool_put(peer_address_text, conn, 100, client_conn_free);

    znx_event_t *e = znx_thread_calloc(sizeof(znx_event_t));
    e->handler = client_timer_handler;

    znx_event_loop_add_timer(e, 10);

    return;
}


void connect_handler(void *data, znx_conn_t *conn, znx_result_t res)
{
    if (res != ZNX_SUCCESS) {
        return;
    }

    conn->read.handler = client_read;
    conn->data = data;

    znx_event_loop_start_io(conn);
    znx_event_loop_enable_read(conn);

    znx_syscall_result_t sres;
    sres = znx_conn_write(conn, "aaaaa", 5);
    if (sres.res != 5) {
        ZNX_DEF_LOG_FATAL("client write failed");
    }
}


static void udp_conn_free(znx_conn_t *conn)
{
    ZNX_DEF_LOG_DEBUG("udp_conn close");
    znx_event_loop_remove_conn(conn);
    znx_conn_close(conn);
    znx_thread_free(conn);
}


static void test_get_and_put_udp_conn()
{
    u_char          buf[128];
    znx_conn_t      *conn;
    znx_str_t       key;
    u_char          *end;
    znx_address_t   address;
    int             port = 0;

    end = znx_snprintf(buf, 512, "UDP-127.0.0.1:8888");
    key.data = buf;
    key.len = (size_t)(end - buf);
    bzero(&address, sizeof(znx_address_t));
    znx_address_parser((u_char *)"127.0.0.1:8888", strlen("127.0.0.1:8888"), &address, &port);
    conn = znx_g_session_pool_get(&key, &address);
    if (conn == NULL) {
        ZNX_DEF_LOG_DEBUG("get udp session %v failed", &key);
    } else {
        ZNX_DEF_LOG_DEBUG("get udp session %v success", &key);
        znx_g_session_pool_put(&key, conn, 200, udp_conn_free);
    }

    end = znx_snprintf(buf, 512, "UDP-127.0.0.1:8889");
    key.data = buf;
    key.len = (size_t)(end - buf);
    bzero(&address, sizeof(znx_address_t));
    znx_address_parser((u_char *)"127.0.0.1:8889", strlen("127.0.0.1:8889"), &address, &port);
    conn = znx_g_session_pool_get(&key, &address);
    if (conn == NULL) {
        ZNX_DEF_LOG_DEBUG("get udp session %v failed", &key);
    } else {
        ZNX_DEF_LOG_DEBUG("get udp session %v success", &key);
        znx_g_session_pool_put(&key, conn, 200, udp_conn_free);
    }

}


void udp_connect_handler(void *data, znx_conn_t *conn, znx_result_t res)
{
    if (res != ZNX_SUCCESS) {
        ZNX_DEF_LOG_ERROR("udp connect failed");
        return;
    }

    znx_str_t *peer_address_text =
        znx_conn_peer_address_text(conn);

    u_char buf[128];
    u_char *end = znx_snprintf(buf, 128, "UDP-%v", peer_address_text);
    znx_str_t key;
    key.data = buf;
    key.len = (size_t)(end - buf);

    znx_g_session_pool_put(&key, conn, 100, udp_conn_free);

    test_get_and_put_udp_conn();
}


znx_bool_t tcp_client_task(void *data)
{
    znx_address_t address;

    bzero(&address, sizeof(znx_address_t));
    if (!znx_address_parser(
        (u_char *)"127.0.0.1:8888", strlen("127.0.0.1:8888"), &address, NULL)) {
        ZNX_DEF_LOG_FATAL("parser %s failed", "127.0.0.1:8888");
    }

    znx_tcp_connect_setup(&address, NULL, 10, NULL, connect_handler, (void *)(intptr_t)1);
    znx_tcp_connect_setup(&address, NULL, 10, NULL, connect_handler, (void *)(intptr_t)1);
    znx_tcp_connect_setup(&address, NULL, 10, NULL, connect_handler, (void *)(intptr_t)1);
    znx_tcp_connect_setup(&address, NULL, 10, NULL, connect_handler, (void *)(intptr_t)1);
    znx_tcp_connect_setup(&address, NULL, 10, NULL, connect_handler, (void *)(intptr_t)1);
    znx_tcp_connect_setup(&address, NULL, 10, NULL, connect_handler, (void *)(intptr_t)1);

    znx_udp_connect_setup(&address, NULL, udp_connect_handler, (void *)(intptr_t)1);
    znx_udp_connect_setup(&address, NULL, udp_connect_handler, (void *)(intptr_t)1);

    bzero(&address, sizeof(znx_address_t));
    if (!znx_address_parser(
        (u_char *)"127.0.0.1:8889", strlen("127.0.0.1:8889"), &address, NULL)) {
        ZNX_DEF_LOG_FATAL("parser %s failed", "127.0.0.1:8889");
    }

    znx_tcp_connect_setup(&address, NULL, 10, NULL, connect_handler, (void *)(intptr_t)2);
    znx_tcp_connect_setup(&address, NULL, 10, NULL, connect_handler, (void *)(intptr_t)2);
    znx_tcp_connect_setup(&address, NULL, 10, NULL, connect_handler, (void *)(intptr_t)2);
    znx_tcp_connect_setup(&address, NULL, 10, NULL, connect_handler, (void *)(intptr_t)2);
    znx_tcp_connect_setup(&address, NULL, 10, NULL, connect_handler, (void *)(intptr_t)2);

    znx_udp_connect_setup(&address, NULL, udp_connect_handler, (void *)(intptr_t)2);
    znx_udp_connect_setup(&address, NULL, udp_connect_handler, (void *)(intptr_t)2);


    return ZNX_TRUE;
}

int main()
{
    znx_srandom((unsigned int)get_cache_time());

    znx_worker_processor_init();

    znx_g_session_pool_init(8);
    znx_worker_processor_add_init_hook(tcp_client_task);

    znx_address_t address;
    bzero(&address, sizeof(znx_address_t));
    if (!znx_address_parser(
        (u_char *)"127.0.0.1:8888", strlen("127.0.0.1:8888"), &address, NULL)) {
        ZNX_DEF_LOG_FATAL("parser %s failed", "127.0.0.1:8888");
    }

    znx_tcp_listener_t *ls = znx_thread_calloc(sizeof(znx_tcp_listener_t));
    znx_tcp_listener_init(ls);
    ls->listen_addr = "127.0.0.1:8888";
    ls->address = address;
    ls->backlog = 1024;
    ls->nodelay = 1;
    ls->reuseaddr = 1;
    ls->accept_handler = tcp_server_accept_handler;

    znx_tcp_listener_register(ls);

    bzero(&address, sizeof(znx_address_t));
    if (!znx_address_parser(
        (u_char *)"127.0.0.1:8889", strlen("127.0.0.1:8889"), &address, NULL)) {
        ZNX_DEF_LOG_FATAL("parser %s failed", "127.0.0.1:8889");
    }

    ls = znx_thread_calloc(sizeof(znx_tcp_listener_t));
    znx_tcp_listener_init(ls);
    ls->listen_addr = "127.0.0.1:8889";
    ls->address = address;
    ls->backlog = 1024;
    ls->nodelay = 1;
    ls->reuseaddr = 1;
    ls->accept_handler = tcp_server_accept_handler;

    znx_tcp_listener_register(ls);

    // open listeners.
    znx_tcp_listeners_open();

    znx_worker_processor_start(1, 1, 1);
    znx_tcp_accepter_start(2);

    usleep(200000);
    proc_exit = ZNX_TRUE;
    while (!znx_worker_processor_all_task_done()) {
        usleep(100000);
    }

    znx_tcp_accepter_stop();
    znx_worker_processor_stop();

    return 0;

}
